# @Version        : 1.0
# @Update Time    : 2025/3/17 21:33
# @File           : anti_spider_middleware.py
# @IDE            : PyCharm
# @Desc           : 防爬虫中间件，用于防止恶意爬虫和请求攻击

import re
import time
from typing import Dict, List, Optional, Set, Pattern, Any
from fastapi import Request, FastAPI
from fastapi.responses import ORJSONResponse
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from src.log import logger


class AntiSpiderMiddleware(BaseHTTPMiddleware):
    """
    防爬虫中间件

    提供以下防护功能：
    - 请求频率限制
    - User-Agent 黑名单检查
    - IP 黑名单管理
    - 可疑请求检测
    """

    def __init__(
        self,
        app: FastAPI,
        rate_limit: int = 100,  # 每分钟最大请求数
        time_window: int = 60,  # 时间窗口（秒）
        user_agent_blacklist: Optional[List[str]] = None,
        storage: Optional[Dict[str, Any]] = None,  # 存储扩展接口
        blacklist_duration: int = 3600,  # 黑名单持续时间（秒）
    ):
        """
        初始化防爬虫中间件

        Args:
            app: FastAPI应用实例
            rate_limit: 每分钟最大请求数
            time_window: 时间窗口（秒）
            user_agent_blacklist: 自定义User-Agent黑名单
            storage: 存储系统接口
            blacklist_duration: 黑名单持续时间（秒）
        """
        super().__init__(app)
        self.rate_limit = rate_limit
        self.time_window = time_window
        self.blacklist_duration = blacklist_duration

        # 初始化默认黑名单模式
        self.user_agent_patterns = [
            r"python-requests.*",
            r"scrapy.*",
            r"curl.*",
            r"wget.*",
            r"^$",  # 空 User-Agent
            r"bot.*",
            r"spider.*",
            r"crawler.*",
            r"headless.*",
            r"phantomjs.*",
            r"selenium.*",
            r"playwright.*",
            r"puppeteer.*",
        ]
        if user_agent_blacklist:
            self.user_agent_patterns.extend(user_agent_blacklist)

        # 编译正则表达式
        self.compiled_patterns = [
            re.compile(p, re.IGNORECASE) for p in self.user_agent_patterns
        ]

        # 初始化存储系统
        self.storage = storage or {
            "request_records": {},
            "blacklist": {},
            "last_cleanup": time.time(),
        }

    async def dispatch(
        self, request: Request, call_next: RequestResponseEndpoint
    ) -> ORJSONResponse:
        """
        中间件调度方法

        Args:
            request: FastAPI请求对象
            call_next: 下一个中间件或路由处理函数

        Returns:
            ORJSONResponse: 响应对象
        """
        try:
            # 获取客户端IP
            client_ip = self._get_client_ip(request)

            # 清理过期数据
            await self._cleanup_expired_data()

            # 黑名单检查
            if await self._is_blacklisted(client_ip):
                logger.warning(f"黑名单IP访问: {client_ip}")
                return self._create_forbidden_response("IP已被封禁")

            # User-Agent检查
            user_agent = request.headers.get("user-agent", "")
            if self._check_suspicious_user_agent(user_agent):
                logger.warning(f"可疑User-Agent: {user_agent} from {client_ip}")
                await self._add_to_blacklist(client_ip)
                return self._create_forbidden_response("可疑的访问行为")

            # 频率限制检查
            if not await self._check_rate_limit(client_ip):
                logger.warning(f"请求频率超限: {client_ip}")
                await self._add_to_blacklist(client_ip)
                return self._create_rate_limit_response()

            # 继续处理请求
            return await call_next(request)

        except Exception as e:
            logger.error(f"防爬虫中间件处理失败: {str(e)}")
            return self._create_forbidden_response("系统异常")

    def _get_client_ip(self, request: Request) -> str:
        """
        获取客户端真实IP

        Args:
            request: FastAPI请求对象

        Returns:
            str: 客户端IP地址
        """
        # 按优先级检查各种IP头
        for header in ["x-forwarded-for", "x-real-ip", "cf-connecting-ip"]:
            if header in request.headers:
                return request.headers[header].split(",")[0].strip()

        # 如果都没有，使用客户端IP
        return request.client.host if request.client else "unknown"

    def _check_suspicious_user_agent(self, user_agent: str) -> bool:
        """
        检查可疑User-Agent

        Args:
            user_agent: User-Agent字符串

        Returns:
            bool: 是否可疑
        """
        return any(pattern.match(user_agent) for pattern in self.compiled_patterns)

    async def _check_rate_limit(self, client_ip: str) -> bool:
        """
        频率限制检查

        Args:
            client_ip: 客户端IP

        Returns:
            bool: 是否通过检查
        """
        now = time.time()
        records = self.storage["request_records"]

        # 获取并清理过期记录
        timestamps = [
            t for t in records.get(client_ip, []) if t > now - self.time_window
        ]

        if len(timestamps) >= self.rate_limit:
            return False

        timestamps.append(now)
        records[client_ip] = timestamps
        return True

    async def _is_blacklisted(self, client_ip: str) -> bool:
        """
        检查IP是否在黑名单中

        Args:
            client_ip: 客户端IP

        Returns:
            bool: 是否在黑名单中
        """
        blacklist = self.storage["blacklist"]
        if client_ip in blacklist:
            # 检查是否过期
            if time.time() - blacklist[client_ip] > self.blacklist_duration:
                del blacklist[client_ip]
                return False
            return True
        return False

    async def _add_to_blacklist(self, client_ip: str) -> None:
        """
        添加IP到黑名单

        Args:
            client_ip: 客户端IP
        """
        self.storage["blacklist"][client_ip] = time.time()

    async def _cleanup_expired_data(self) -> None:
        """清理过期数据"""
        now = time.time()
        # 每5分钟清理一次
        if now - self.storage["last_cleanup"] < 300:
            return

        # 清理过期的请求记录
        records = self.storage["request_records"]
        for ip in list(records.keys()):
            records[ip] = [t for t in records[ip] if t > now - self.time_window]
            if not records[ip]:
                del records[ip]

        # 清理过期的黑名单记录
        blacklist = self.storage["blacklist"]
        for ip in list(blacklist.keys()):
            if now - blacklist[ip] > self.blacklist_duration:
                del blacklist[ip]

        self.storage["last_cleanup"] = now

    def _create_forbidden_response(self, detail: str) -> ORJSONResponse:
        """
        创建禁止访问响应

        Args:
            detail: 错误详情

        Returns:
            ORJSONResponse: 响应对象
        """
        return ORJSONResponse(status_code=403, content={"detail": detail})

    def _create_rate_limit_response(self) -> ORJSONResponse:
        """
        创建频率限制响应

        Returns:
            ORJSONResponse: 响应对象
        """
        return ORJSONResponse(
            status_code=429,
            content={"detail": "请求过于频繁，请稍后再试"},
            headers={"Retry-After": str(self.time_window)},
        )
